cmake_minimum_required(VERSION 3.20.0 FATAL_ERROR)

if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
  set(CMAKE_CUDA_ARCHITECTURES 70)
endif()
project(AgradTest CXX CUDA)

# set(CMAKE_BUILD_TYPE Debug)
# set(CMAKE_BUILD_TYPE Release)

set(CXX_FLAGS_DEBUG -O0 -g -G --compiler-options -Wall
                    -Xcompiler "-fopenmp -rdynamic" 
                    -ftz=true --expt-extended-lambda)
                    
set(CXX_FLAGS_RELEASE -O3 --compiler-options -Wall 
                      -Xptxas -v
                      --expt-extended-lambda)



# add_subdirectory(GSWITCH)

find_package(
  Torch REQUIRED
  # PATHS .pyenv/lib/python3.8/site-packages/torch
  # NO_DEFAULT_PATH
)

#find openMP
find_package(OpenMP)
find_package(Python COMPONENTS Development)

if(OPENMP_FOUND)
  message("OpenMP found with flags\n"
          "  C_FLAGS: ${OpenMP_C_FLAGS}\n"
          "  CXX_FLAGS: ${OpenMP_CXX_FLAGS}\n"
          "  CUDA_FLAGS: ${OpenMP_CUDA_FLAGS}")
  set(OpenMP_CUDA_FLAGS -Xcompiler ${OpenMP_CXX_FLAGS})
endif()

if(TORCH_FOUND)
  message("Torch FOUND")
else()
  message("PyTorch not found!")
endif()


add_subdirectory(torch_grad_test)

add_subdirectory(torch_bench)
